package com.wg.core.id;

import cn.hutool.core.lang.func.VoidFunc1;
import cn.hutool.core.util.StrUtil;
import cn.hutool.db.Db;
import cn.hutool.db.Entity;
import cn.hutool.db.transaction.TransactionLevel;
import cn.hutool.log.Log;
import cn.hutool.log.LogFactory;
import com.wg.db.ds.DataSourceConfig;

import javax.sql.DataSource;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 唯一ID生成器
 * <p>
 * 特性：
 * 1. 分段获取ID，减少DB压力
 * 2. 支持自适应步长扩展
 * 3. 支持最大步长上限保护
 * 4. 长时间无大需求时自动收缩步长
 */
public class WgIdGenerator {
    private static final Log log = LogFactory.get();

    /**
     * 默认步长
     */
    private static final long DEFAULT_STEP = 1000;
    /**
     * 最大步长上限
     */
    private static final long MAX_STEP = 100_000;
    /**
     * 收缩阈值：10分钟
     */
    private static final long SHRINK_IDLE_MS = 10 * 60 * 1000;

    private static final String CREATE_TABLE_SQL = "CREATE TABLE IF NOT EXISTS  `use_id_alloc`  (" +
            "`key` char(128) CHARACTER SET utf8 COLLATE utf8_general_ci NOT NULL," +
            "`max_id` bigint NULL DEFAULT 1000000," +
            "`step` bigint NULL DEFAULT 1000," +
            "`description` varchar(255) CHARACTER SET utf8 COLLATE utf8_general_ci NULL DEFAULT NULL COMMENT '描述'," +
            "`update_time` timestamp NULL DEFAULT NULL ON UPDATE CURRENT_TIMESTAMP," +
            "PRIMARY KEY (`key`) USING BTREE) " +
            "ENGINE = InnoDB CHARACTER SET = utf8 COLLATE = utf8_general_ci ROW_FORMAT = Dynamic;";

    private final Map<String, IdCache> idCacheMap = new ConcurrentHashMap<>(16);
    private final Db useDb;

    public WgIdGenerator(DataSourceConfig config) {
        try {
            DataSource dataSource = config.buildDataSource();
            this.useDb = Db.use(dataSource);
            init();
        } catch (Exception e) {
            throw new RuntimeException("Init WgIdGenerator failed", e);
        }
    }

    public WgIdGenerator(Db db) {
        this.useDb = db;
        init();
    }

    /**
     * 获取下一个唯一ID
     */
    public long nextId(String key) throws Exception {
        if (StrUtil.isEmpty(key)) {
            throw new IllegalArgumentException("key is null");
        }

        IdCache idCache = idCacheMap.computeIfAbsent(key, IdCache::new);
        synchronized (idCache) {
            if (idCache.isNeedUpdate()) {
                reloadFromDb(key, idCache, 1); // 至少需要1个
            }
            return idCache.next();
        }
    }

    /**
     * 批量获取一段唯一ID
     *
     * @param key  业务key
     * @param size 需要的数量
     * @return long[2] => [start, end]
     */
    public long[] nextRange(String key, int size) throws Exception {
        if (StrUtil.isEmpty(key)) {
            throw new IllegalArgumentException("key is null");
        }
        if (size <= 0) {
            throw new IllegalArgumentException("size must > 0");
        }

        IdCache idCache = idCacheMap.computeIfAbsent(key, IdCache::new);
        synchronized (idCache) {
            if (idCache.remaining() < size) {
                reloadFromDb(key, idCache, size);
                if (idCache.remaining() < size) {
                    throw new RuntimeException("ID allocation failed, size=" + size);
                }
            }

            long start = idCache.getMinId() + 1;
            long end = idCache.getMinId() + size;
            idCache.advance(size);

            return new long[]{start, end};
        }
    }

    /**
     * 事务内从数据库加载新的区间
     */
    private void reloadFromDb(String key, IdCache cache, long requiredStep) throws Exception {
        useDb.tx(TransactionLevel.SERIALIZABLE, (VoidFunc1<Db>) parameter -> {
            // 保证key存在
            parameter.execute("INSERT IGNORE INTO `use_id_alloc`(`key`) VALUES (?);", key);

            // 查询当前 step
            Entity current = parameter.queryOne("SELECT max_id, step FROM use_id_alloc WHERE `key` = ?;", key);
            long dbStep = current.get("step", DEFAULT_STEP);

            // 计算最终步长
            long finalStep = Math.max(dbStep, requiredStep);

            // 收缩逻辑：长时间无大需求，恢复默认
            long now = System.currentTimeMillis();
            if (finalStep > DEFAULT_STEP
                    && (now - cache.getLastUpdateTime() > SHRINK_IDLE_MS)
                    && requiredStep < dbStep / 2) {
                log.info("Key={} shrink step from {} -> {} due to idle time", key, finalStep, DEFAULT_STEP);
                finalStep = DEFAULT_STEP;
            }

            // 上限保护
            if (finalStep > MAX_STEP) {
                log.warn("Key={} requiredStep={} exceeds MAX_STEP={}, force limit.", key, requiredStep, MAX_STEP);
                finalStep = MAX_STEP;
            }

            // 更新 max_id & step
            parameter.execute("UPDATE use_id_alloc SET `max_id` = `max_id` + ?, `step` = ? WHERE `key` = ?;",
                    finalStep, finalStep, key);

            // 查询最新区间
            Entity entity = parameter.queryOne("SELECT max_id, step FROM use_id_alloc WHERE `key` = ?;", key);
            if (entity != null) {
                long maxId = entity.getLong("max_id");
                long step = entity.getLong("step");
                cache.updateRange(maxId, step);

                log.debug("Reload ID range: key={} minId={} maxId={} step={}", key, cache.getMinId(), cache.getMaxId(), step);
            }
        });
    }

    /**
     * 初始化表 & 预加载已有的key
     */
    private void init() {
        try {
            useDb.execute(CREATE_TABLE_SQL);
            List<Entity> entities = useDb.query("SELECT `key`, max_id, step FROM use_id_alloc");
            for (Entity entity : entities) {
                String key = entity.getStr("key");
                long maxId = entity.get("max_id", 0L);
                long step = entity.get("step", DEFAULT_STEP);

                IdCache cache = new IdCache(key);
                // 初始化时，默认让它下一次请求时触发 reload
                cache.updateRange(maxId, 0);
                cache.setStep(step);

                idCacheMap.put(key, cache);
            }
        } catch (Exception e) {
            throw new RuntimeException("Init WgIdGenerator failed", e);
        }
    }

    /**
     * ID 缓存
     */
    private static class IdCache {
        private long minId;
        private long maxId;
        private long lastUpdateTime; // 最近一次更新的时间戳

        public IdCache(String key) {
            this.lastUpdateTime = System.currentTimeMillis();
        }

        public boolean isNeedUpdate() {
            return minId >= maxId;
        }

        public long next() {
            return ++minId;
        }

        public void updateRange(long newMaxId, long step) {
            this.maxId = newMaxId;
            this.minId = newMaxId - step;
            this.lastUpdateTime = System.currentTimeMillis();
        }

        public void setStep(long step) {
        }

        public long getMinId() {
            return minId;
        }

        public long getMaxId() {
            return maxId;
        }

        public long remaining() {
            return maxId - minId;
        }

        public void advance(long size) {
            this.minId += size;
        }

        public long getLastUpdateTime() {
            return lastUpdateTime;
        }
    }
}
